package com.octopus.core;

import com.octopus.core.collector.Collector;
import com.octopus.core.collector.DefaultCollector;
import com.octopus.core.downloader.DefaultDownloader;
import com.octopus.core.downloader.DownloadException;
import com.octopus.core.downloader.Downloader;
import com.octopus.core.downloader.proxy.DefaultProxyProvider;
import com.octopus.core.downloader.proxy.HostBindProxyProvider;
import com.octopus.core.downloader.proxy.Proxy;
import com.octopus.core.downloader.proxy.ProxyProvider;
import com.octopus.core.downloader.proxy.URLRegexProxyProvider;
import com.octopus.core.limiter.DefaultRateLimitPlanner;
import com.octopus.core.limiter.RateLimitPlanner;
import com.octopus.core.limiter.RateLimiter;
import com.octopus.core.matcher.RegexMatcher;
import com.octopus.core.parser.DefaultTextParser;
import com.octopus.core.parser.ParseException;
import com.octopus.core.parser.ParseResult;
import com.octopus.core.parser.Parser;
import com.octopus.core.scheduler.DefaultScheduler;
import com.octopus.core.scheduler.Scheduler;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Pattern;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;

/**
 * 爬虫核心入口，负责管理整个爬虫工作流程和生命周期管理
 *
 * @author shoulai.yang@gmail.com
 * @date 2019/07/31
 */
@Slf4j
public final class Octopus {

  /** 默认爬虫线程数 */
  private static final int DEFAULT_THREAD_NUM = 5;

  /** 爬虫监听器集合 */
  private List<OctopusListener> listeners = new ArrayList<>();

  /** HTTP 下载器 */
  private Downloader downloader;

  /** 下载请求调度器 */
  private Scheduler scheduler;

  /** 页面解析器 */
  private Parser parser;

  /** 数据收集 */
  private Collector collector;

  /** 爬虫线程数目 */
  private int threads = DEFAULT_THREAD_NUM;

  /** 当前状态引擎状态 */
  private AtomicReference<State> state = new AtomicReference<>(State.INIT);

  /** 速率限制器 */
  private RateLimitPlanner rateLimitPlanner;

  /** 爬虫信号量 */
  private Semaphore workersSemaphore;

  /** 爬虫线程池 */
  private ExecutorService workersPool = null;

  /** 代理信息 */
  private ProxyProvider proxyProvider;

  /** 没有任务时是否关闭引擎 */
  private boolean autoStop = false;

  /** 引擎空闲状态监听 */
  private final Object idleWatcher = new Object();

  public static class Builder {

    private List<OctopusListener> listeners = new ArrayList<>();

    private RateLimitPlanner rateLimitPlanner;

    private Downloader downloader = new DefaultDownloader();

    private Scheduler scheduler = new DefaultScheduler();

    private Parser parser = new DefaultTextParser();

    private Collector collector = new DefaultCollector();

    private int threads = DEFAULT_THREAD_NUM;

    private HostBindProxyProvider hostBindProxyProvider = new HostBindProxyProvider();

    private URLRegexProxyProvider urlRegexProxyProvider = new URLRegexProxyProvider();

    private ProxyProvider defaultProxyProvider;

    private boolean autoStop;

    public Builder downloader(@NonNull Downloader downloader) {
      this.downloader = downloader;
      return this;
    }

    public Builder autoStop() {
      this.autoStop = true;
      return this;
    }

    public Builder scheduler(@NonNull Scheduler scheduler) {
      this.scheduler = scheduler;
      return this;
    }

    public Builder parser(@NonNull Parser parser) {
      this.parser = parser;
      return this;
    }

    public Builder collector(@NonNull Collector collector) {
      this.collector = collector;
      return this;
    }

    public Builder listener(@NonNull OctopusListener listener) {
      this.listeners.add(listener);
      return this;
    }

    public Builder rateLimitPlanner(@NonNull RateLimitPlanner rateLimitPlanner) {
      this.rateLimitPlanner = rateLimitPlanner;
      return this;
    }

    public Builder rateLimit(int num, @NonNull TimeUnit unit) {
      return this.rateLimitPlanner(new DefaultRateLimitPlanner(new RateLimiter(num, unit)));
    }

    public Builder proxy(@NonNull String bindHost, @NonNull Proxy proxy) {
      this.hostBindProxyProvider.addProxy(bindHost, proxy);
      return this;
    }

    public Builder proxy(@NonNull Pattern pattern, @NonNull Proxy proxy) {
      this.urlRegexProxyProvider.addProxy(RegexMatcher.of(pattern), proxy);
      return this;
    }

    public Builder proxy(@NonNull Proxy proxy) {
      this.defaultProxyProvider = new DefaultProxyProvider(proxy);
      return this;
    }

    public Builder proxyProvider(ProxyProvider proxyProvider) {
      this.defaultProxyProvider = proxyProvider;
      return this;
    }

    public Builder threads(int threads) {
      if (threads <= 0) {
        log.warn("thead number must greater than 0, use default instead");
        this.threads = DEFAULT_THREAD_NUM;
      } else {
        this.threads = threads;
      }
      return this;
    }

    public Octopus build() {
      Octopus octopus = new Octopus();
      octopus.collector = this.collector;
      octopus.downloader = this.downloader;
      octopus.parser = this.parser;
      octopus.threads = this.threads;
      octopus.scheduler = this.scheduler;
      if (this.rateLimitPlanner != null) {
        octopus.rateLimitPlanner = this.rateLimitPlanner;
        octopus.listeners.add(this.rateLimitPlanner);
      }
      this.listeners.add(this.collector);
      this.listeners.add(this.downloader);
      this.listeners.add(this.scheduler);
      octopus.listeners.addAll(this.listeners);
      octopus.workersSemaphore = new Semaphore(octopus.threads);
      octopus.proxyProvider =
          request -> {
            Proxy proxy = Builder.this.hostBindProxyProvider.provide(request);
            if (proxy == null) {
              proxy = Builder.this.urlRegexProxyProvider.provide(request);
              if (proxy == null && Builder.this.defaultProxyProvider != null) {
                proxy = Builder.this.defaultProxyProvider.provide(request);
              }
            }
            return proxy;
          };
      octopus.autoStop = this.autoStop;
      return octopus;
    }
  }

  /** 引擎状态 */
  enum State {
    /** 初始化 */
    INIT,
    /** 运行 */
    RUNNING,
    /** 空闲 */
    IDLE,
    /** 停止 */
    STOPPED
  }

  static class WorkersThreadFactory implements ThreadFactory {

    private int workerSequence = 1;

    @Override
    public Thread newThread(Runnable r) {
      Thread thread = new Thread(r);
      thread.setName("[SPIDER-" + workerSequence++ + "]");
      return thread;
    }
  }

  private Octopus() {}

  /**
   * 添加爬取任务
   *
   * @param requests 任务列表
   */
  public void addRequest(@NonNull Request... requests) {
    if (this.state.get() != State.RUNNING && this.state.get() != State.IDLE) {
      throw new IllegalStateException("Octopus in illegal state " + this.state.get().toString());
    }
    for (Request request : requests) {
      if (!request.repeatable() && this.scheduler.visited(request) && request.retryTimes() == 0) {
        log.info("Repeated request! {}", request);
        continue;
      }
      if (!this.scheduler.add(request)) {
        log.error("Request can not put into scheduler. {}", request);
      } else {
        this.scheduler.visit(request);
      }
    }
    if (this.state.get() == State.IDLE) {
      synchronized (idleWatcher) {
        this.idleWatcher.notifyAll();
      }
    }
  }

  /**
   * 开启引擎
   *
   * @param urls 种子链接
   */
  public void start(@NonNull String... urls) {
    Request[] requests = new Request[urls.length];
    for (int i = 0; i < urls.length; i++) {
      requests[i] = Request.me(urls[i]);
    }
    this.start(requests);
  }

  /**
   * 开启引擎
   *
   * @param requests 种子任务
   */
  public void start(Request... requests) {
    if (this.state.get() == State.RUNNING || this.state.get() == State.IDLE) {
      throw new IllegalStateException("Octopus in illegal state " + this.state.get().toString());
    }
    // init workers thread pool
    this.setupThreadPool();
    // starting
    this.setState(State.RUNNING);
    this.onStart();
    this.addRequest(requests);
    while (!Thread.currentThread().isInterrupted()) {
      State state = this.state.get();
      if (state == State.IDLE) {
        try {
          synchronized (idleWatcher) {
            this.idleWatcher.wait();
          }
        } catch (InterruptedException e) {
          e.printStackTrace();
        }
        if (this.state.get() != State.STOPPED) {
          this.setState(State.RUNNING);
        }
        continue;
      }
      if (state != State.RUNNING) {
        break;
      }
      if (this.scheduler.remaining() > 0) {
        final Request request = this.scheduler.get();
        if (request == null) {
          continue;
        }
        try {
          this.workersSemaphore.acquire();
          this.workersPool.execute(
              () -> {
                try {
                  if (this.rateLimitPlanner != null) {
                    RateLimiter limiter = this.rateLimitPlanner.getRateLimiter(request);
                    if (limiter != null) {
                      limiter.acquire();
                    }
                  }
                  Proxy proxy =
                      this.proxyProvider == null ? null : this.proxyProvider.provide(request);
                  log.debug(
                      "Download {} {}, remaining {}",
                      request,
                      proxy == null ? "" : "via " + proxy,
                      this.scheduler.remaining());
                  Response response = this.downloader.download(request, proxy);
                  this.onDownloaded(response);
                  if (response.httpStatus() == 200) {
                    try {
                      ParseResult result = this.parser.parse(response);
                      result.requests().forEach(this::addRequest);
                      if (result.body() != null && !result.body().isEmpty()) {
                        this.collector.collect(result.body(), response);
                      }
                    } catch (ParseException e) {
                      this.onException(e);
                    }
                  }
                } catch (Exception e) {
                  if (e instanceof DownloadException) {
                    if (request.retriable()) {
                      log.info("Retry! {}", request);
                      request.fail();
                      this.addRequest(request);
                    } else {
                      log.error(String.format("Download error! %s", request), e);
                      this.onException(e);
                    }
                  } else {
                    log.error("Octopus run error!", e);
                    this.onException(e);
                  }
                } finally {
                  this.workersSemaphore.release();
                }
              });
        } catch (InterruptedException e) {
          e.printStackTrace();
        }

      } else {
        // stop if no request found and no worker running
        if (this.scheduler.remaining() == 0
            && this.workersSemaphore.availablePermits() == this.threads) {
          if (this.autoStop) {
            log.info("No more request found, octopus will stop.");
            this.stop();
          } else {
            log.info("No more request found, octopus will idle.");
            this.setState(State.IDLE);
            this.onIdle();
          }
        }
      }
    }
  }

  private void setupThreadPool() {
    this.workersPool =
        new ThreadPoolExecutor(
            this.threads,
            this.threads,
            5,
            TimeUnit.SECONDS,
            new LinkedBlockingQueue<>(),
            new WorkersThreadFactory());
  }

  private void setState(State state) {
    for (; ; ) {
      if (this.state.compareAndSet(this.state.get(), state)) {
        break;
      }
    }
  }

  /** 停止引擎 */
  public void stop() {
    if (this.state.get() != State.RUNNING && this.state.get() != State.IDLE) {
      throw new IllegalStateException("Octopus in illegal state " + this.state.get().toString());
    }
    this.setState(State.STOPPED);
    synchronized (idleWatcher) {
      this.idleWatcher.notifyAll();
    }
    this.workersPool.shutdown();
    this.onStop();
  }

  private void onStart() {
    log.info("Octopus started!");
    this.listeners.forEach(OctopusListener::onStart);
  }

  private void onIdle() {
    log.info("Octopus idle!");
    this.listeners.forEach(OctopusListener::onIdle);
  }

  private void onStop() {
    log.debug("Notify all listeners that octopus will stop.");
    this.listeners.forEach(OctopusListener::onStop);
    log.info("Octopus stopped!");
  }

  private void onException(Exception e) {
    this.listeners.forEach(l -> l.onException(e));
  }

  private void onDownloaded(Response response) {
    this.listeners.forEach(l -> l.onDownloaded(response));
  }
}
